#!/bin/bash

# Quick SSH Test Script for 5GC Testbed
# This script verifies SSH key authentication for all VMs

set -e

# VM Configuration
VM1="192.168.1.206"  # Core Network
VM2="192.168.1.219"  # UPF
VM3="192.168.1.221"  # gNodeB
VM4="192.168.1.205"  # UE

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

success() {
    echo -e "✓ ${GREEN}$1${NC}"
}

error() {
    echo -e "✗ ${RED}$1${NC}"
}

info() {
    echo -e "${BLUE}$1${NC}"
}

# Function to test SSH connectivity
test_ssh() {
    local vm=$1
    local vm_name=$2
    local user=${3:-ubuntu}
    
    echo -n "Testing $vm_name ($user@$vm)... "
    
    if ssh -o ConnectTimeout=10 -o BatchMode=yes $user@$vm "echo 'SSH connection successful'" 2>/dev/null; then
        success "SSH key authentication works"
        return 0
    else
        error "SSH key authentication failed"
        return 1
    fi
}

# Function to test basic system info
test_system_info() {
    local vm=$1
    local vm_name=$2
    local user=${3:-ubuntu}
    
    echo -n "Getting system info for $vm_name... "
    
    if system_info=$(ssh -o ConnectTimeout=10 -o BatchMode=yes $user@$vm "hostname && whoami && uname -a" 2>/dev/null); then
        success "System info retrieved"
        echo "  Hostname: $(echo "$system_info" | head -1)"
        echo "  User: $(echo "$system_info" | sed -n '2p')"
        echo "  System: $(echo "$system_info" | tail -1)"
    else
        error "Failed to get system info"
        return 1
    fi
}

# Main test function
main() {
    echo
    echo "=========================================="
    echo "      Quick SSH Test for 5GC Testbed"
    echo "=========================================="
    echo "Date: $(date)"
    echo "=========================================="
    echo
    
    log "Testing SSH key authentication for all VMs..."
    echo
    
    # Test SSH connectivity
    vms=("$VM1:VM1 (Core)" "$VM2:VM2 (UPF)" "$VM3:VM3 (gNB)" "$VM4:VM4 (UE)")
    
    success_count=0
    total_count=0
    
    for vm_info in "${vms[@]}"; do
        vm_ip=$(echo $vm_info | cut -d: -f1)
        vm_name=$(echo $vm_info | cut -d: -f2)
        
        echo "Testing $vm_name ($vm_ip)..."
        
        if test_ssh $vm_ip $vm_name; then
            test_system_info $vm_ip $vm_name
            ((success_count++))
        fi
        
        ((total_count++))
        echo
    done
    
    # Summary
    echo "=========================================="
    echo "           Test Summary"
    echo "=========================================="
    echo "Total VMs: $total_count"
    echo -e "Successful: ${GREEN}$success_count${NC}"
    echo -e "Failed: ${RED}$((total_count - success_count))${NC}"
    echo "=========================================="
    
    if [ $success_count -eq $total_count ]; then
        echo
        success "All VMs are accessible via SSH key authentication!"
        info "You can now run the 5GC scripts:"
        info "  ./scripts/health-check.sh"
        info "  ./scripts/startup-5gc.sh"
        info "  ./tests/run-all-tests.sh"
    else
        echo
        error "Some VMs failed SSH authentication!"
        info "Please check:"
        info "  1. SSH keys are properly copied to all VMs"
        info "  2. SSH service is running on all VMs"
        info "  3. Network connectivity is established"
    fi
}

# Handle command line arguments
case "${1:-}" in
    "--help"|"-h")
        echo "Usage: $0 [OPTIONS]"
        echo
        echo "Options:"
        echo "  -h, --help          Show this help message"
        echo
        echo "This script will:"
        echo "  1. Test SSH key authentication to all VMs"
        echo "  2. Get basic system information from each VM"
        echo "  3. Provide a summary of connectivity status"
        ;;
    *)
        main
        ;;
esac 